#include <vector>
#include<iostream>
#include<stdio.h>
#include <algorithm>
#include "log_sum_exp_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"
using namespace std;

// void print_arr(int32_t* array, int32_t arr_len){
//     for(int32_t i=0; i<arr_len; ++i)
//     {
//         printf("%d ", array[i]);
        
//     }
//     printf("\n");
// }

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    TilingData tiling;
    uint32_t sizeofdatatype;
    uint32_t totalLengthAligned;

    // 1. 获取平台信息
    uint64_t ub_size;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto socVersion = ascendcPlatform.GetSocVersion();
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);

    // 2. 获取数据信息
    const bool* pkeep_dims = context->GetAttrs()->GetBool(0);
    const gert::TypedContinuousVector<int64_t> *axes2 = 
                        context->GetAttrs()->GetListInt(0);

    uint32_t dataNum;
    uint32_t totalLength = context->GetInputTensor(0)->GetShapeSize();

    auto shape = context->GetInputTensor(0)->GetOriginShape();
    const int32_t dimNum = shape.GetDimNum();
    const int32_t axesLen = axes2->GetSize();
    
    
    int32_t reduce_num = axesLen;
    int32_t reduce_shapes[reduce_num];
    int32_t reduce_strides[reduce_num];
    int32_t reduce_strides2[reduce_num];

    int32_t output_num = dimNum-axesLen;
    int32_t output_shapes[output_num];
    int32_t output_strides[output_num];
    int32_t output_strides2[output_num];

    int32_t outputLength = 1;
    int32_t inputLength = totalLength;

    vector<int32_t> axes(axesLen);
    int32_t input_shape[dimNum];
    int32_t input_strides[dimNum];
    int64_t axis;
    for(int32_t i = 0; i < axesLen; i++){
        axis = *(axes2->GetData()+i);
        if(axis < 0){
            axis += dimNum;
        }
        axes[i] = int32_t(axis);
    }

    sort(axes.begin(), axes.end());
    // auto last = std::unique(axes.begin(), axes.end());
    // axes.erase(last, axes.end());

    for(int32_t i = 0; i < dimNum; i++){
        input_shape[i] = shape.GetDim(i);
    }

    input_strides[dimNum-1] = 1;
    for(int32_t i = dimNum-2; i >=0; i--){
        input_strides[i] = input_strides[i+1]*input_shape[i+1];
    }

    int32_t reduce_index = 0;
    int32_t output_index = 0;
    for(int32_t i = 0; i < dimNum; i++){
        if(i==axes[reduce_index])
        {
            reduce_shapes[reduce_index] = input_shape[i];
            reduce_strides[reduce_index] = input_strides[i];
            reduce_index+=1;
        }
        else
        {
            output_shapes[output_index] = input_shape[i];
            output_strides[output_index] = input_strides[i];
            outputLength = outputLength*output_shapes[output_index];
            output_index+=1;
        }
    }


    output_strides2[output_num-1] = 1;
    for(int32_t i = output_num-2; i >=0; i--){
        output_strides2[i] = output_strides2[i+1]*output_shapes[i+1];
    }

    reduce_strides2[reduce_num-1] = 1;
    for(int32_t i = reduce_num-2; i >=0; i--){
        reduce_strides2[i] = reduce_strides2[i+1]*reduce_shapes[i+1];
    }


    
    tiling.set_reduce_shapes(reduce_shapes);
    tiling.set_output_shapes(output_shapes);
    tiling.set_reduce_strides(reduce_strides);
    tiling.set_reduce_strides2(reduce_strides2);
    tiling.set_output_strides(output_strides);
    tiling.set_output_strides2(output_strides2);

    tiling.set_reduce_num(reduce_num);
    tiling.set_output_num(output_num);
    tiling.set_inputLength(inputLength);
    tiling.set_outputLength(outputLength);

    context->SetBlockDim(1);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
                        context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t* currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    return ge::GRAPH_SUCCESS;
}
}

namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
static graphStatus InferDataType(gert::InferDataTypeContext* context)
{
    const auto inputDataType = context->GetInputDataType(0);
    context->SetOutputDataType(0, inputDataType);
    return ge::GRAPH_SUCCESS;
}
}

namespace ops {
class LogSumExp : public OpDef {
public:
    explicit LogSumExp(const char* name) : OpDef(name)
    {
        this->Input("x")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        // this->Input("axes")
        //         .ParamType(REQUIRED)
        //         .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
        //         .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
        //         .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

        this->Output("y")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("dim").AttrType(OPTIONAL).ListInt(vector<int64_t>(1));
        this->Attr("keep_dims").AttrType(OPTIONAL).Bool(false);
        

        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);
        this->AICore()
            .SetTiling(optiling::TilingFunc)
            .AddConfig("ascend310b");
    }
};
OP_ADD(LogSumExp);
}
